package cn.sciento.websocket.handler;


import cn.sciento.core.convertor.ApplicationContextHelper;
import cn.sciento.core.oauth.CustomTokenConverter;
import cn.sciento.core.oauth.CustomUserDetails;
import cn.sciento.websocket.config.WebSocketConfig;
import cn.sciento.websocket.helper.SocketMessageHandler;
import cn.sciento.websocket.redis.BrokerSessionRedis;
import cn.sciento.websocket.redis.SessionUserRedis;
import cn.sciento.websocket.registry.SessionRegistry;
import cn.sciento.websocket.vo.MsgVO;
import cn.sciento.websocket.vo.UserVO;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.commons.lang3.reflect.FieldUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.ResponseEntity;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.util.MultiValueMap;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.AbstractWebSocketHandler;

import java.io.IOException;
import java.util.Map;
import java.util.Objects;

public class WebSocketHandler extends AbstractWebSocketHandler {
    private static final Logger logger = LoggerFactory.getLogger(WebSocketHandler.class);

    private RestTemplate restTemplate = (RestTemplate) ApplicationContextHelper.getContext().getBean("restTemplate");

    private ObjectMapper objectMapper = ApplicationContextHelper.getContext().getBean(ObjectMapper.class);

    private WebSocketConfig config;

    public WebSocketHandler(WebSocketConfig config) {
        this.config = config;
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        try {
            String sessionId = null;
            if (session instanceof org.springframework.web.socket.adapter.standard.StandardWebSocketSession) {
                sessionId = session.getId();
            } else if (session instanceof org.springframework.web.socket.sockjs.transport.session.WebSocketServerSockJsSession) {
                sessionId = ((WebSocketSession) FieldUtils.readField(session, "webSocketSession", true)).getId();
            } else {
                session.close();
            }
            String token = String.valueOf(session.getAttributes().get("access_token"));
            CustomUserDetails customUserDetails = getAuthentication(token);
            if (customUserDetails != null) {
                Long tenantId = customUserDetails.getTenantId();
                Long userId = customUserDetails.getUserId();
                Long roleId = customUserDetails.getRoleId();
                String brokerId = SessionRegistry.getBrokerId();
                UserVO user = new UserVO(sessionId, tenantId, roleId, token, brokerId);
                logger.debug("connection success. userId : {}", userId);
                SessionRegistry.addSession(session, sessionId, userId);
                BrokerSessionRedis.refreshCache(brokerId, userId, user);
                SessionUserRedis.refreshCache(user);
            } else {
                session.close();
            }
        } catch (Exception e) {
            try {
                session.close();
            } catch (IOException ex) {
                logger.debug("session : {} closed failed.", session);
            }
            logger.warn("webSocket connection failed. message : {}, token : {}", e.getMessage(), session.getAttributes().get("access_token"));
        }
    }

    private CustomUserDetails getAuthentication(String token) {
        HttpHeaders headers = new HttpHeaders();
        headers.set("Authorization", "bearer " + token);
        HttpEntity<String> entity = new HttpEntity("parameters", (MultiValueMap) headers);
        ResponseEntity<Map<String, Object>> responseEntity = this.restTemplate.exchange(this.config.getOauthUrl() + "?access_token=" + token, HttpMethod.GET, entity, new ParameterizedTypeReference<Map<String, Object>>() {

        });
        Map<String, Object> result = (Map<String, Object>) responseEntity.getBody();
        CustomTokenConverter customTokenConverter = new CustomTokenConverter();
        OAuth2Authentication oAuth2Authentication = customTokenConverter.extractAuthentication(result);
        Object object = oAuth2Authentication.getDetails();
        if (object instanceof CustomUserDetails) {
            return (CustomUserDetails) object;
        }
        return null;
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) {
        if (message instanceof org.springframework.web.socket.TextMessage) {
            String messageStr = String.valueOf(message.getPayload());
            if (Objects.equals(messageStr, this.config.getHeartbeat())) {
                return;
            }
            Map<String, SocketMessageHandler> beansOfType = ApplicationContextHelper.getContext().getBeansOfType(SocketMessageHandler.class);
            beansOfType.forEach((name, bean) -> {
                MsgVO msg;
                if (bean.needPrincipal()) {
                    String token = (String) session.getAttributes().get("access_token");
                    bean.setCustomUserDetails(getAuthentication(token));
                }
                try {
                    msg = this.objectMapper.readValue(messageStr, MsgVO.class);
                } catch (Exception e) {
                    logger.error("handleMessage error , message : {}", message.getPayload());
                    return;
                }
                bean.processMessage(msg);
            });
        } else if (message instanceof org.springframework.web.socket.BinaryMessage) {
            logger.info("Unexpected WebSocket message type: {}", message);
        } else if (message instanceof org.springframework.web.socket.PongMessage) {
            logger.info("Unexpected WebSocket message type: {}", message);
        } else {
            logger.info("Unexpected WebSocket message type: {}", message);
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) {
        clearSession(session);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        clearSession(session);
    }

    private void clearSession(WebSocketSession session) {
        String sessionId;
        try {
            if (session instanceof org.springframework.web.socket.adapter.standard.StandardWebSocketSession) {
                sessionId = session.getId();
            } else if (session instanceof org.springframework.web.socket.sockjs.transport.session.WebSocketServerSockJsSession) {
                sessionId = ((WebSocketSession) FieldUtils.readField(session, "webSocketSession", true)).getId();
            } else {
                return;
            }
        } catch (Exception e) {
            logger.warn("webSocket disConnection failed.");
            return;
        }
        Long userId = SessionRegistry.getUser(sessionId);
        if (userId != null) {
            String brokerId = SessionRegistry.getBrokerId();
            UserVO invalidUser = BrokerSessionRedis.deleteCache(brokerId, userId, sessionId);
            SessionUserRedis.deleteCache(invalidUser);
            SessionRegistry.removeSession(sessionId);
        }
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }
}
